
import warnings
warnings.filterwarnings('ignore')
import pandas as pd
import numpy as np
from numpy import abs
from numpy import log
from numpy import sign
import sys 
sys.path.append('code')
from feature_eng.alpha_series.alpha_utility import industry_dict


def correlation(x,y,n):
    cor_lst = []
    for i in range(x.shape[1]):
        cor_i = (x.iloc[:,i]).rolling(n, min_periods=1).corr(y.iloc[:,i])
        cor_lst.append(cor_i)
    corr = pd.concat(cor_lst, axis=1)
    corr.replace([np.inf, -np.inf], np.nan, inplace = True)
    corr = corr.fillna(0)
    return corr

def covariance(x,y,n):
    cor_lst = []
    for i in range(x.shape[1]):
        cor_i = (x.iloc[:,i]).rolling(n, min_periods=1).cov(y.iloc[:,i])
        cor_lst.append(cor_i)
    corr = pd.concat(cor_lst, axis=1)
    corr.replace([np.inf, -np.inf], np.nan, inplace = True)
    corr = corr.fillna(0)
    return corr   


def indneutralize(df):
    lagret_v = []
    ind_dic = industry_dict()
    for i in df.columns:
        for k, v in ind_dic.items():
            if i in v:
                v_i = df[i] - df[v].mean(axis=1)
                lagret_v.append(v_i)
    lagRet_vol = pd.concat(lagret_v, axis=1)
    lagRet_vol.columns = df.columns
    dc = lagRet_vol
    dc.replace([np.inf, -np.inf], np.nan, inplace = True)
    dc = dc.fillna(0)
    return dc

def ts_sum(df, window=10):
    """
    Wrapper function to estimate rolling sum.
    :param df: a pandas DataFrame.
    :param window: the rolling window.
    :return: a pandas DataFrame with the time-series min over the past 'window' days.
    """
    
    return df.rolling(window, min_periods=1).sum()

def Mean(df, window=10):
    """
    Wrapper function to estimate Mean.
    :param df: a pandas DataFrame.
    :param window: the rolling window.
    :return: a pandas DataFrame with the time-series min over the past 'window' days.
    """
    return df.rolling(window, min_periods=1).mean()

def stddev(df, window=10):
    """
    Wrapper function to estimate rolling standard deviation.
    :param df: a pandas DataFrame.
    :param window: the rolling window.
    :return: a pandas DataFrame with the time-series min over the past 'window' days.
    """
    return df.rolling(window, min_periods=1).std()

def ts_rank(df, window=10):
    """
    Wrapper function to estimate rolling rank.
    :param df: a pandas DataFrame.
    :param window: the rolling window.
    :return: a pandas DataFrame with the time-series rank over the past window days.
    """
    return df.rolling(window, min_periods=1).rank(pct=True)


def product(df, window=10):
    """
    Wrapper function to estimate rolling product.
    :param df: a pandas DataFrame.
    :param window: the rolling window.
    :return: a pandas DataFrame with the time-series product over the past 'window' days.
    """
    return df.rolling(window, min_periods=1).apply(lambda x: x.sum() , raw=True, engine="numba")

def ts_min(df, window=10):
    """
    Wrapper function to estimate rolling min.
    :param df: a pandas DataFrame.
    :param window: the rolling window.
    :return: a pandas DataFrame with the time-series min over the past 'window' days.
    """
    return df.rolling(window, min_periods=1).min()

def ts_max(df, window=10):
    """
    Wrapper function to estimate rolling min.
    :param df: a pandas DataFrame.
    :param window: the rolling window.
    :return: a pandas DataFrame with the time-series max over the past 'window' days.
    """
    return df.rolling(window, min_periods=1).max()

def delta(df, period=1):
    """
    Wrapper function to estimate difference.
    :param df: a pandas DataFrame.
    :param period: the difference grade.
    :return: a pandas DataFrame with today’s value minus the value 'period' days ago.
    """
    return df.diff(period)

def delay(df, period=1):
    """
    Wrapper function to estimate lag.
    :param df: a pandas DataFrame.
    :param period: the lag grade.
    :return: a pandas DataFrame with lagged time series
    """
    return df.shift(period)

def rank(df):
    """
    Cross sectional rank
    :param df: a pandas DataFrame.
    :return: a pandas DataFrame with rank along columns.
    """
    #return df.rank(axis=1, pct=True)
    df_rank = df.rank(pct=True, axis=1)
    return df_rank

def scale(df, k=1):
    """
    Scaling time serie.
    :param df: a pandas DataFrame.
    :param k: scaling factor.
    :return: a pandas DataFrame rescaled df such that sum(abs(df)) = k
    """
    # df.mul(k).div(np.abs(df).sum())
    return df/abs(df).expanding(min_periods=1).sum()

def gen_y(x):
    weights = np.arange(1,len(x)+1,1)
    y = weights / weights.sum()   
    return (x*y).sum()

def decay_linear(df, window=10):
    """
    Linear weighted moving average implementation.
    :param df: a pandas DataFrame.
    :param period: the LWMA period
    :return: a pandas DataFrame with the LWMA.
    """
      #y是和为1的权重
    lst = []
    for i in range(df.shape[1]):
        sr = df.iloc[:,i]
        sr1 = sr.rolling(window, min_periods=1).apply(gen_y, raw=True, engine="numba")
        lst.append(sr1)
    dc = pd.concat(lst, axis=1)
    dc.replace([np.inf, -np.inf], np.nan, inplace = True)
    dc = dc.fillna(0)
    return dc


def gen_wma(x):
    weights = 0.9*np.arange(len(x)-1,0-1,-1)
    return (x*weights).sum()

def Wma(df, window=10):
    lst = []
    for i in range(df.shape[1]):
        sr = df.iloc[:,i]
        sr1 = sr.rolling(window, min_periods=1).apply(gen_wma, raw=True, engine="numba")
        lst.append(sr1)
    dc = pd.concat(lst, axis=1)
    dc.replace([np.inf, -np.inf], np.nan, inplace = True)
    dc = dc.fillna(0)
    return dc

def ts_argmin(df, window=10):
    lst = []
    for i in range(df.shape[1]):
        sr = df.iloc[:,i]
        sr1 = sr.rolling(window, min_periods=1).apply(np.argmin, raw=True, engine="numba") + 1
        lst.append(sr1)
    dc = pd.concat(lst, axis=1)
    dc.replace([np.inf, -np.inf], np.nan, inplace = True)
    dc = dc.fillna(0)
    return dc

def ts_argmax(df, window=10):
    lst = []
    for i in range(df.shape[1]):
        sr = df.iloc[:,i]
        sr1 = sr.rolling(window, min_periods=1).apply(np.argmax, raw=True, engine="numba") + 1
        lst.append(sr1)
    dc = pd.concat(lst, axis=1)
    dc.replace([np.inf, -np.inf], np.nan, inplace = True)
    dc = dc.fillna(0)
    return dc

def Count(part,window):
    #计算前n期满足条件condition的样本个数,此时输入的part为0、1变量
    return part.rolling(window, min_periods=1, closed='left').sum()




def Max(sr,n):
    #比较取大函数
    return np.maximum(n, sr)

def Max2(sr1,sr2):
    #比较取小函数
    return np.maximum(sr1, sr2)

def Min(sr,n):
    #比较取小函数
    return np.minimum(n, sr)

def Min2(sr1,sr2):
    #比较取小函数
    return np.minimum(sr1, sr2)

def ewma(sr,n,m):
    return sr.ewm(alpha=m/n, adjust=False).mean()

def Sumif(part,window):
    #对前n项条件求和，part为条件筛选后的数据
    return part.rolling(window, min_periods=1, closed='left').sum()

def Lowday(df,window):
    #计算sr前window期时间序列中最小值距离当前时点的间隔
    lst = []
    for i in range(df.shape[1]):
        sr = df.iloc[:,i]
        sr1 = sr.rolling(window, min_periods=1, closed='left').apply(np.argmin,raw=True, engine="numba") + 1
        lst.append(sr1)
    dc = pd.concat(lst, axis=1)
    dc.replace([np.inf, -np.inf], np.nan, inplace = True)
    dc = dc.fillna(0)
    return window - dc


def Highday(df,window):
    #计算sr前window期时间序列中最大值距离当前时点的间隔
    lst = []
    for i in range(df.shape[1]):
        sr = df.iloc[:,i]
        sr1 = sr.rolling(window, min_periods=1, closed='left').apply(np.argmax,raw=True, engine="numba") + 1
        lst.append(sr1)
    dc = pd.concat(lst, axis=1)
    dc.replace([np.inf, -np.inf], np.nan, inplace = True)
    dc = dc.fillna(0)
    return window - dc

